#!/bin/bash

###################################################################
# @Author:             SongShaohua                                #
# @Date:               2021-03-31 14:40:32                        #
# @Last Modified by:   Updated                                    #
# @Last Modified time: 2025-09-29 00:20:46                        #
###################################################################

# 定义当前目录和日志文件
CURRENT_DIR=$(cd "$(dirname $0)";pwd)
LOG_FILE="${CURRENT_DIR}/autoexssh_$(date +%Y%m%d%H%M%S).log"

# 默认参数
USER=""
PORT="22"
KEY_FILE=""
PASSWORD=""
CONFIG_FILE=""
USE_CONFIG=false
DEBUG=false

# SSH选项配置
SSH_OPTIONS="-o StrictHostKeyChecking=no -o ConnectTimeout=30"
SCP_OPTIONS="-o StrictHostKeyChecking=no -o ConnectTimeout=30"

# 日志颜色配置
ERROR_COLOR='\033[1;91m'
WARNING_COLOR='\033[1;93m'
INFO_COLOR='\033[0;94m'
MESSAGE_COLOR='\033[1;92m'
NC='\033[0m'

# 主机信息存储
declare -A HOST_INFO
HOST_LIST=()

# 日志记录函数
function log() {
    local level=$1
    local message=$2
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")
    local log_entry="[${timestamp}] [${level}] ${message}"
    
    # 写入日志文件
    echo "${log_entry}" >> "${LOG_FILE}"
    
    # 根据日志级别和DEBUG模式决定是否输出到控制台
    case "${level}" in
        ERROR) echo -e "${ERROR_COLOR}${log_entry}${NC}" ;;
        MESSAGE) 
            if [[ "${message}" == *"开始验证主机连通性"* ]] || \
               [[ "${message}" == *"主机连通性验证完成"* ]] || \
               [[ "${message}" == *"成功读取"*"个主机配置"* ]] || \
               [[ "${message}" == *"开始执行SSH密钥自动配置脚本"* ]] || \
               [[ "${message}" == *"SSH密钥自动配置脚本执行完成"* ]] || \
               "$DEBUG" == true; then
                echo -e "${MESSAGE_COLOR}${log_entry}${NC}"
            fi
            ;;
        WARNING) 
            if [ "$DEBUG" == true ]; then
                echo -e "${WARNING_COLOR}${log_entry}${NC}"
            fi
            ;;
        DEBUG) 
            if [ "$DEBUG" == true ]; then
                echo -e "${INFO_COLOR}${log_entry}${NC}"
            fi
            ;;
    esac
}

# 帮助信息
function help() {
    echo "说明: "
    echo "  一个用于自动配置 ssh 互信的工具"
    echo "  用法一适用于所有主机密码相同的情况，只需要配置 all_hosts 主机地址即可"
    echo "  用法二适用于所有主机密码相同或不同的情况,需要严格按照配置格式配置"
    echo "  用法一: sh $0 -u <用户名> -P <密码> [-p <端口>] [-f <密钥文件>] [-d]"
    echo "  用法二: sh $0 -c <配置文件> [-f <密钥文件>] [-d]"
    echo "选项:"
    echo "  -u 服务器用户 (当不使用 -c 参数时必需)"
    echo "  -P 服务器密码 (当不使用 -c 参数时必需)"
    echo "  -p 服务器端口 (当不使用 -c 参数时可选，默认22)"
    echo "  -c 指定配置文件路径，格式: ip:port:user:password 每行一个"
    echo "  -f 指定密钥文件路径 (可选，默认使用 ssh-keygen 默认路径: ~/.ssh/)"
    echo "  -d 开启debug模式，显示详细输出"
    echo "示例:"
    echo "  sh $0 -u root -P password -p 22 -f ~/.ssh/my_rsa_key"
    echo "  sh $0 -c hosts.conf -f ~/.ssh/my_rsa_key -d"
    exit 99
}

# 密码加密和解密函数
function encrypt_password() {
    local password=$1
    echo "${password}" | base64
}

function decrypt_password() {
    local encrypted=$1
    echo "${encrypted}" | base64 -d
}

# 加密配置文件中的密码
function encrypt_config_passwords() {
    local config_file=$1
    local temp_file="${config_file}.tmp"
    
    log "MESSAGE" "开始加密配置文件中的密码"
    
    while IFS=':' read -r ip port user password || [ -n "$ip" ]; do
        if [ -n "$password" ]; then
            encrypted=$(encrypt_password "$password")
            echo "${ip}:${port}:${user}:${encrypted}" >> "${temp_file}"
        else
            echo "${ip}:${port}:${user}:" >> "${temp_file}"
        fi
    done < "${config_file}"
    
    mv "${temp_file}" "${config_file}"
    log "MESSAGE" "配置文件密码加密完成，加密后的密码已写入 ${config_file}"
}

# 解析命令行参数
function parse_args() {
    if [ $# -eq 0 ]; then
        help
    fi

    while getopts "u:P:p:f:c:d?" opt; do
        case $opt in
            u) USER=$OPTARG ;;
            P) PASSWORD=$OPTARG ;;
            p) PORT=$OPTARG ;;
            f) KEY_FILE=$OPTARG ;;
            c) CONFIG_FILE=$OPTARG; USE_CONFIG=true ;;
            d) DEBUG=true ;;
            ?) help ;;
            *) help ;;
        esac
    done

    # 根据模式读取主机配置
    if [ "$USE_CONFIG" = true ]; then
        if [ -z "$CONFIG_FILE" ] || [ ! -f "$CONFIG_FILE" ]; then
            log "ERROR" "配置文件不存在: $CONFIG_FILE"
            help
        fi
        read_config_file
        validate_hosts
    else
        if [ -z "$USER" ] || [ -z "$PASSWORD" ]; then
            log "ERROR" "缺少必需参数 -u 和 -P"
            help
        fi
        read_all_hosts
        validate_hosts
    fi
    
    # 检查是否有有效主机
    if [ ${#HOST_LIST[@]} -eq 0 ]; then
        log "ERROR" "没有有效的主机配置"
        exit 1
    fi
    
    log "MESSAGE" "成功读取 ${#HOST_LIST[@]} 个主机配置"
}

# 从配置文件读取主机信息
function read_config_file() {
    log "MESSAGE" "从配置文件读取主机信息: $CONFIG_FILE"
    
    local line_num=0
    while IFS=':' read -r ip port user password || [ -n "$ip" ]; do
        line_num=$((line_num + 1))
        
        # 跳过空行和注释
        if [[ "$ip" =~ ^\s*# ]] || [ -z "$ip" ]; then
            continue
        fi
        
        # 验证必填字段
        if [ -z "$ip" ]; then
            log "ERROR" "配置文件第${line_num}行IP地址不能为空"
            exit 1
        fi
        
        if [ -z "$port" ]; then
            log "ERROR" "配置文件第${line_num}行端口不能为空"
            exit 1
        fi
        
        if [ -z "$user" ]; then
            log "ERROR" "配置文件第${line_num}行用户名不能为空"
            exit 1
        fi
        
        if [ -z "$password" ]; then
            log "ERROR" "配置文件第${line_num}行密码不能为空"
            exit 1
        fi
        
        # 存储主机信息
        HOST_INFO["${ip}_ip"]="${ip}"
        HOST_INFO["${ip}_port"]="${port}"
        HOST_INFO["${ip}_user"]="${user}"
        HOST_INFO["${ip}_password"]="${password}"
        HOST_LIST+=("${ip}")
        
        log "DEBUG" "读取主机配置: IP=${ip}, 端口=${port}, 用户=${user}"
    done < "$CONFIG_FILE"
}

# 从 all_hosts 读取主机信息
function read_all_hosts() {
    if [ ! -f "${CURRENT_DIR}/all_hosts" ]; then
        log "ERROR" "all_hosts 文件不存在: ${CURRENT_DIR}/all_hosts"
        exit 1
    fi
    
    log "MESSAGE" "从 all_hosts 文件读取主机信息"
    
    local line_num=0
    while read -r ip || [ -n "$ip" ]; do
        line_num=$((line_num + 1))
        
        # 跳过空行和注释
        if [[ "$ip" =~ ^\s*# ]] || [ -z "$ip" ]; then
            continue
        fi
        
        # 存储主机信息
        HOST_INFO["${ip}_ip"]="${ip}"
        HOST_INFO["${ip}_port"]="${PORT}"
        HOST_INFO["${ip}_user"]="${USER}"
        HOST_INFO["${ip}_password"]="${PASSWORD}"
        HOST_LIST+=("${ip}")
        
        log "DEBUG" "读取主机配置: IP=${ip}, 端口=${PORT}, 用户=${USER}"
    done < "${CURRENT_DIR}/all_hosts"
    
    # 替换日志中的明文密码
    if [ -n "$PASSWORD" ]; then
        local encrypted=$(encrypt_password "$PASSWORD")
        sed -i "s/${PASSWORD}/${encrypted}/g" "$LOG_FILE"
        log "MESSAGE" "日志中的明文密码已加密处理"
    fi
}

# 验证主机连通性
function validate_hosts() {
    log "MESSAGE" "开始验证主机连通性"
    
    for ip in "${HOST_LIST[@]}"; do
        local port=${HOST_INFO["${ip}_port"]}
        
        # Ping测试
        if command -v ping &> /dev/null; then
            if [ "$DEBUG" == true ]; then
                log "DEBUG" "正在ping主机 ${ip}..."
            fi
            ping -c 1 -W 2 "$ip" &> /dev/null
            if [ $? -ne 0 ]; then
                log "ERROR" "主机 ${ip} 无法ping通,请检查防火墙或路由信息"
                exit 1
            fi
        fi
        
        # 端口验证
        if command -v nc &> /dev/null; then
            if [ "$DEBUG" == true ]; then
                log "DEBUG" "正在测试主机 ${ip}:${port} 的连通性..."
            fi
            nc -z -w 3 "$ip" "$port" &> /dev/null
            if [ $? -ne 0 ]; then
                log "ERROR" "主机 ${ip}:${port} 端口不可达,请检查防火墙或路由信息"
                exit 1
            fi
        else
            # 使用bash内置TCP连接
            if [ "$DEBUG" == true ]; then
                log "DEBUG" "nc命令不可用，使用bash内置TCP连接能力测试主机 ${ip}:${port} 的连通性..."
            fi
            timeout 3 bash -c "exec 3<>/dev/tcp/$ip/$port && exec 3>&- && exec 3<&-" &> /dev/null
            if [ $? -ne 0 ]; then
                log "ERROR" "主机 ${ip}:${port} 端口不可达,请检查防火墙或路由信息"
                exit 1
            fi
        fi
    done
    
    log "MESSAGE" "主机连通性验证完成"
}

# 获取密钥文件路径
function get_key_path() {
    if [ -z "$KEY_FILE" ]; then
        echo "~/.ssh/id_rsa"
    else
        echo "$KEY_FILE"
    fi
}

# 在主机上生成密钥对
function generate_key_on_host() {
    local host=$1
    local user=${HOST_INFO["${host}_user"]}
    local port=${HOST_INFO["${host}_port"]}
    local password=${HOST_INFO["${host}_password"]}
    local actual_key_path=$(get_key_path)
    
    log "MESSAGE" "在主机 $host 生成密钥对"
    
    # 确保.ssh目录存在
    /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
    spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} mkdir -p ~/.ssh
    expect {
        "*yes/no*" { send "yes\r";exp_continue }
        "*password*" { send "${password}\r";exp_continue}
        eof
    }
EOF
    
    # 生成密钥对
    if [ -z "$KEY_FILE" ]; then
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
    spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} ssh-keygen -t rsa -N ''
    expect {
        "*yes/no*" { send "yes\r";exp_continue }
        "*password*" { send "${password}\r";exp_continue}
        "*Enter file in which to save the key*" { send "\r";exp_continue }
        "Overwrite*" { send "y\r";exp_continue}
        eof
    }
EOF
    else
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} ssh-keygen -t rsa -N '' -f ${KEY_FILE}
            expect {
                "*yes/no*" { send "yes\r";exp_continue }
                "*password*" { send "${password}\r";exp_continue}
                "Overwrite*" { send "y\r";exp_continue}
                eof
            }
EOF
    fi
    
    # 检查密钥生成成功
    /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
        spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} ls -la ${actual_key_path}
        expect {
            "*yes/no*" { send "yes\r";exp_continue }
            "*password*" { send "${password}\r";exp_continue}
            "No such file or directory" { exit 1 }
            eof
        }
EOF
    
    if [ $? -eq 0 ]; then
        log "MESSAGE" "主机 ${host} 密钥对生成成功: ${actual_key_path}"
        return 0
    else
        log "ERROR" "主机 ${host} 密钥对生成失败"
        return 1
    fi
}

# 收集所有节点公钥
function collect_all_public_keys() {
    local actual_key_path=$(get_key_path)
    mkdir -p ${CURRENT_DIR}/temp_keys
    > ${CURRENT_DIR}/temp_keys/authorized_keys
    
    log "MESSAGE" "开始收集所有节点的公钥"
    
    for host in "${HOST_LIST[@]}"; do
        local user=${HOST_INFO["${host}_user"]}
        local port=${HOST_INFO["${host}_port"]}
        local password=${HOST_INFO["${host}_password"]}
        
        log "MESSAGE" "从主机 ${host} 获取公钥"
        
        # 获取公钥
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn scp ${SCP_OPTIONS} -p -P ${port} ${user}@${host}:${actual_key_path}.pub ${CURRENT_DIR}/temp_keys/${host}.pub
            expect {
                "*yes/no*" { send "yes\r";exp_continue}
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
        
        # 添加到authorized_keys
        if [ -f ${CURRENT_DIR}/temp_keys/${host}.pub ]; then
            cat ${CURRENT_DIR}/temp_keys/${host}.pub >> ${CURRENT_DIR}/temp_keys/authorized_keys
            log "MESSAGE" "已获取 ${host} 的公钥并添加到authorized_keys"
        else
            log "WARNING" "无法获取 ${host} 的公钥"
        fi
    done
    
    log "MESSAGE" "所有节点公钥收集完成"
}

# 分发authorized_keys到所有节点
function distribute_authorized_keys() {
    local actual_key_path=$(get_key_path)
    
    for host in "${HOST_LIST[@]}"; do
        local user=${HOST_INFO["${host}_user"]}
        local port=${HOST_INFO["${host}_port"]}
        local password=${HOST_INFO["${host}_password"]}
        
        log "MESSAGE" "分发合并后的authorized_keys到: $host"
        
        # 确保.ssh目录存在并设置权限
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} "mkdir -p ~/.ssh && chmod 700 ~/.ssh"
            expect {
                "*yes/no*" { send "yes\r";exp_continue }
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
        
        # 分发authorized_keys
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn scp ${SCP_OPTIONS} -p -P ${port} ${CURRENT_DIR}/temp_keys/authorized_keys ${user}@${host}:~/.ssh/
            expect {
                "*yes/no*" { send "yes\r";exp_continue}
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
        
        # 设置权限
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} chmod 600 ~/.ssh/authorized_keys
            expect {
                "*yes/no*" { send "yes\r";exp_continue }
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
    done
    
    log "MESSAGE" "authorized_keys 分发完成"
}

# 复制密钥文件到所有节点
function copy_keys_to_all_nodes() {
    local actual_key_path=$(get_key_path)
    local source_host=${HOST_LIST[0]}
    local source_user=${HOST_INFO["${source_host}_user"]}
    local source_port=${HOST_INFO["${source_host}_port"]}
    local source_password=${HOST_INFO["${source_host}_password"]}
    
    log "MESSAGE" "开始复制密钥文件到所有节点"
    
    for host in "${HOST_LIST[@]}"; do
        # 跳过源节点
        if [ "$host" == "$source_host" ]; then
            continue
        fi
        
        local user=${HOST_INFO["${host}_user"]}
        local port=${HOST_INFO["${host}_port"]}
        local password=${HOST_INFO["${host}_password"]}
        
        log "MESSAGE" "将密钥复制到 ${host}"
        
        # 复制私钥
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn scp ${SCP_OPTIONS} -p -P ${source_port} ${source_user}@${source_host}:${actual_key_path} ${user}@${host}:${actual_key_path}
            expect {
                "*yes/no*" { send "yes\r";exp_continue}
                "*password*" { send "${source_password}\r";exp_continue}
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
        
        # 复制公钥
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn scp ${SCP_OPTIONS} -p -P ${source_port} ${source_user}@${source_host}:${actual_key_path}.pub ${user}@${host}:${actual_key_path}.pub
            expect {
                "*yes/no*" { send "yes\r";exp_continue}
                "*password*" { send "${source_password}\r";exp_continue}
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
        
        # 设置权限
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} "chmod 600 ${actual_key_path} && chmod 644 ${actual_key_path}.pub"
            expect {
                "*yes/no*" { send "yes\r";exp_continue }
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
    done
    
    log "MESSAGE" "密钥文件复制完成"
}

# 测试互信连接
function test_ssh_connection() {
    local actual_key_path=$(get_key_path)
    log "MESSAGE" "开始测试节点间互信连接"
    
    # 测试每个节点到其他节点的连接
    for source_host in "${HOST_LIST[@]}"; do
        local source_user=${HOST_INFO["${source_host}_user"]}
        local source_port=${HOST_INFO["${source_host}_port"]}
        
        for target_host in "${HOST_LIST[@]}"; do
            # 跳过自己连接自己
            if [ "$source_host" == "$target_host" ]; then
                continue
            fi
            
            local target_user=${HOST_INFO["${target_host}_user"]}
            local target_port=${HOST_INFO["${target_host}_port"]}
            
            log "MESSAGE" "测试从 $source_host 到 $target_host 的SSH连接"
            
            # 使用密钥文件连接测试
            /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
                spawn ssh ${SSH_OPTIONS} -i ${actual_key_path} -p ${source_port} ${source_user}@${source_host} ssh ${SSH_OPTIONS} -i ${actual_key_path} -p ${target_port} ${target_user}@${target_host} date
                expect {
                    "*yes/no*" { send "yes\r";exp_continue }
                    "*password*" { exit 1 }
                    "Permission denied" { exit 1 }
                    eof
                }
EOF
            
            if [ $? -eq 0 ]; then
                log "MESSAGE" "从 $source_host 到 $target_host 的SSH连接成功"
            else
                log "ERROR" "从 $source_host 到 $target_host 的SSH连接失败"
            fi
        done
    done
    
    # 输出测试命令示例
    if [ -n "$KEY_FILE" ]; then
        local first_host=${HOST_LIST[0]}
        local second_host=${HOST_LIST[1]:-${HOST_LIST[0]}}
        local first_user=${HOST_INFO["${first_host}_user"]}
        local first_port=${HOST_INFO["${first_host}_port"]}
        local second_user=${HOST_INFO["${second_host}_user"]}
        local second_port=${HOST_INFO["${second_host}_port"]}
        
        log "MESSAGE" "连接测试命令示例:"
        log "MESSAGE" "ssh -i ${KEY_FILE} -p ${first_port} ${first_user}@${first_host} ssh -i ${KEY_FILE} -p ${second_port} ${second_user}@${second_host} date"
    fi
}

# 清理临时文件
function cleanup() {
    log "MESSAGE" "清理临时文件"
    rm -rf ${CURRENT_DIR}/temp_keys
    log "MESSAGE" "临时文件清理完成"
}

# 确保密钥权限正确设置
function ensure_node_permissions() {
    local actual_key_path=$(get_key_path)
    
    for host in "${HOST_LIST[@]}"; do
        local user=${HOST_INFO["${host}_user"]}
        local port=${HOST_INFO["${host}_port"]}
        local password=${HOST_INFO["${host}_password"]}
        
        log "MESSAGE" "确保节点 ${host} 密钥权限正确设置"
        
        /usr/bin/expect <<EOF >>$LOG_FILE 2>&1
            spawn ssh ${SSH_OPTIONS} ${user}@${host} -p ${port} "chmod 600 ${actual_key_path} && chmod 644 ${actual_key_path}.pub && chmod 700 ~/.ssh"
            expect {
                "*yes/no*" { send "yes\r";exp_continue }
                "*password*" { send "${password}\r";exp_continue}
                eof
            }
EOF
    done
}

# 主函数
function main() {
    parse_args "$@"
    log "MESSAGE" "开始执行SSH密钥自动配置脚本"

    # 生成密钥对
    for host in "${HOST_LIST[@]}"; do
        if ! generate_key_on_host "$host"; then
            log "ERROR" "在主机 ${host} 生成密钥失败，终止操作"
            exit 1
        fi
    done
    
    # 设置权限
    ensure_node_permissions
    
    # 收集和分发公钥
    collect_all_public_keys
    distribute_authorized_keys
    
    # 测试互信连接
    test_ssh_connection
    
    # 清理临时文件
    cleanup
    
    # 加密配置文件密码
    if [ "$USE_CONFIG" = true ] && [ -f "$CONFIG_FILE" ]; then
        encrypt_config_passwords "$CONFIG_FILE"
    fi
    
    log "MESSAGE" "SSH密钥自动配置脚本执行完成，共处理 ${#HOST_LIST[@]} 个主机"
}

# 执行主函数
main "$@"
